from langchain_chroma import Chroma
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

from base import embeddings, model

#3. 使用 Chroma 构建向量库 并把向量库存储到 指定目录
vector_store_dir = "store/redis_util"
#pip install chromadb
vectorStore = Chroma(
    persist_directory=vector_store_dir,
    embedding_function= embeddings
)

#检索器
retriever = vectorStore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "k":5,
        "score_threshold":0.5
    }
)

#构建提示模板
prompt = PromptTemplate.from_template(
    """
    你是一个JAVA只能开发助手，你可以完全使用“向量数据库的内容”回答问题,
    如果找不到相关资料直接回答；‘无法知晓该内容’\n\n
    向量数据库的内容：{context}
    问题：{question}
    回答：
    """
)

chain= ({"context":retriever,"question":RunnablePassthrough()}
        | prompt
        | model
        )
resp = chain.invoke("RedisUtil缓存双删的用法")
#保存聊天的历史记录
# chatHistory = ChatMessageHistory()
# def get_session_history():
#     return chatHistory
#
# doMessage = RunnableWithMessageHistory(
#     chain,
#     get_session_history=get_session_history,
# )
#
# resp = doMessage.invoke("RedisUtil缓存双删的用法")

print(resp.content)

